{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b970b917-333d-4059-bf33-ac128fa3bd0c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "from functools import partial\n",
    "\n",
    "from tqdm import tqdm\n",
    "import mindspore\n",
    "from mindspore import nn, ops, Tensor\n",
    "\n",
    "from mindnlp import load_dataset\n",
    "from mindnlp.transforms import BasicTokenizer, PadTransform, Lookup, AddToken, Truncate\n",
    "from mindnlp.vocab import Vocab\n",
    "from mindnlp.modules import StaticGRU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebe0cd3a-9c54-48b3-9220-26ecc4e05408",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "split = ['train', 'validation', 'test']\n",
    "ds_dict = load_dataset('bentrevett/multi30k', split=split)\n",
    "train_dataset, valid_dataset, test_dataset = ds_dict['train'], ds_dict['validation'], ds_dict['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "218e21fd-5b39-4167-921c-93437257dbdd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenizer = BasicTokenizer(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b78bcadd-6b87-48de-83ea-fec28624a1a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.map(tokenizer, 'en').map(tokenizer, 'de')\n",
    "valid_dataset = valid_dataset.map(tokenizer, 'en').map(tokenizer, 'de')\n",
    "test_dataset = test_dataset.map(tokenizer, 'en').map(tokenizer, 'de')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "015f9038-d6d7-46bc-813a-cda3c374965b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for de, en in test_dataset:\n",
    "    print(f'de = {de}')\n",
    "    print(f'en = {en}')\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7888a19d-aa3a-489c-a72f-9e173a49f16c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "de_vocab = Vocab.from_dataset(train_dataset, 'de', freq_range=(2, None), \\\n",
    "                              special_tokens=['<unk>', '<pad>', '<bos>', '<eos>'], special_first=True)\n",
    "en_vocab = Vocab.from_dataset(train_dataset, 'en', freq_range=(2, None), \\\n",
    "                              special_tokens=['<unk>', '<pad>', '<bos>', '<eos>'], special_first=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec7c5936-97d4-4b03-a88c-1d1e13ff3bc3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "max_len = 32\n",
    "\n",
    "begin_add = AddToken('<bos>')\n",
    "end_add = AddToken('<eos>', False)\n",
    "truncate = Truncate(max_len - 2)\n",
    "\n",
    "de_lookup_op = Lookup(de_vocab, unk_token='<unk>')\n",
    "de_pad_op = PadTransform(max_len, de_vocab('<pad>'), return_length=True)\n",
    "\n",
    "en_lookup_op = Lookup(en_vocab, unk_token='<unk>')\n",
    "en_pad_op = PadTransform(max_len, en_vocab('<pad>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d0763bc-5de9-4fe9-a26d-b7509025df3e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.map([truncate, begin_add, end_add, de_lookup_op, de_pad_op], 'de', ['de', 'de_len']) \\\n",
    "                            .map([truncate, begin_add, end_add, en_lookup_op, en_pad_op], 'en')\n",
    "\n",
    "valid_dataset = valid_dataset.map([truncate, begin_add, end_add, de_lookup_op, de_pad_op], 'de', ['de', 'de_len']) \\\n",
    "                            .map([truncate, begin_add, end_add, en_lookup_op, en_pad_op], 'en')\n",
    "\n",
    "test_dataset = test_dataset.map([truncate, begin_add, end_add, de_lookup_op, de_pad_op], 'de', ['de', 'de_len']) \\\n",
    "                            .map([truncate, begin_add, end_add, en_lookup_op, en_pad_op], 'en')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1a6c5d3-7992-471c-8954-108c942b1641",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "train_dataset = train_dataset.batch(batch_size, drop_remainder=True)\n",
    "valid_dataset = valid_dataset.batch(batch_size)\n",
    "test_dataset = test_dataset.batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f394b27-fcf9-4b5c-a88e-931959bca250",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Cell):\n",
    "    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(input_dim, emb_dim)  # Embedding层\n",
    "        self.rnn = StaticGRU(emb_dim, enc_hid_dim, bidirectional=True)  # 双向GRU层\n",
    "        self.fc = nn.Dense(enc_hid_dim * 2, dec_hid_dim)  # 全连接层\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout)  # dropout，防止过拟合\n",
    "\n",
    "    def construct(self, src, src_len):\n",
    "        \"\"\"构建编码器\n",
    "\n",
    "        Args:\n",
    "            src: 源序列，为已经转换为数字索引并统一长度的序列；shape = [src len, batch_size]\n",
    "            src_len: 有效长度；shape = [batch_size, ]\n",
    "        \"\"\"\n",
    "        # 将输入源序列转化为向量，并进行暂退（dropout）\n",
    "        # shape = [src len, batch size, emb dim]\n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        # 计算输出\n",
    "        # shape = [src len, batch size, enc hid dim*2]\n",
    "        outputs, hidden = self.rnn(embedded)\n",
    "        # 为适配解码器，合并两个上下文函数\n",
    "        # shape = [batch size, dec hid dim]\n",
    "        hidden = ops.tanh(self.fc(ops.concat((hidden[-2, :, :], hidden[-1, :, :]), axis=1)))\n",
    "\n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d9b881-af36-4866-8c3a-42985630a73c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Attention(nn.Cell):\n",
    "    def __init__(self, enc_hid_dim, dec_hid_dim):\n",
    "        super().__init__()\n",
    "        # attention线性层\n",
    "        self.attn = nn.Dense((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)\n",
    "        # v， 用不带有bias的线性层表示\n",
    "        # shape = [1, dec hid dim]\n",
    "        self.v = nn.Dense(dec_hid_dim, 1, has_bias=False)\n",
    "\n",
    "    def construct(self, hidden, encoder_outputs, mask):\n",
    "        \"\"\"Attention层\n",
    "\n",
    "        Args:\n",
    "            hidden: 解码器上一个时刻的隐藏状态；shape = [batch size, dec hid dim]\n",
    "            encoder_outputs: 编码器的输出，前向与反向RNN的隐藏状态；shape = [src len, batch size, enc hid dim * 2]\n",
    "            mask: 将<pad>占位符的注意力权重替换为0或者很小的数值；shape = [batch size, src len]\n",
    "        \"\"\"\n",
    "        src_len = encoder_outputs.shape[0]\n",
    "\n",
    "        # 重复解码器隐藏状态src len次，对齐维度\n",
    "        # shape = [batch size, src len, dec hid dim]\n",
    "        hidden = ops.tile(hidden.expand_dims(1), (1, src_len, 1))\n",
    "\n",
    "        # 将编码器输出中的第1、2维度进行交换，对齐维度\n",
    "        # shape = [batch size, src len, enc hid dim*2]\n",
    "        encoder_outputs = encoder_outputs.transpose(1, 0, 2)\n",
    "\n",
    "        # 计算E_t\n",
    "        # shape = [batch size, src len, dec hid dim]\n",
    "        energy = ops.tanh(self.attn(ops.concat((hidden, encoder_outputs), axis=2)))\n",
    "\n",
    "        # 计算v * E_t\n",
    "        # shape = [batch size, src len]\n",
    "        attention = self.v(energy).squeeze(2)\n",
    "\n",
    "        # 不需要考虑序列中<pad>占位符的注意力权重\n",
    "        attention = attention.masked_fill(mask == 0, -1e10)\n",
    "\n",
    "        return ops.softmax(attention, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cff0e45e-73af-4f07-adc6-334d555eceed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Cell):\n",
    "    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):\n",
    "        super().__init__()\n",
    "        self.output_dim = output_dim\n",
    "        self.attention = attention\n",
    "\n",
    "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        self.rnn = StaticGRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)\n",
    "        self.fc_out = nn.Dense((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def construct(self, inputs, hidden, encoder_outputs, mask):\n",
    "        \"\"\"构建解码器\n",
    "\n",
    "        Args:\n",
    "            input: 输入的单词；shape = [batch size]\n",
    "            hidden: 解码器上一时刻的隐藏状态；shape = [batch size, dec hid dim]\n",
    "            encoder_outputs: 编码器的输出，前向与反向RNN的隐藏状态；shape = [src len, batch size, enc hid dim * 2]\n",
    "            mask: 将<pad>占位符的注意力权重替换为0或者很小的数值；shape = [batch size, src len]\n",
    "        \"\"\"\n",
    "\n",
    "        # 为输入增加额外维度\n",
    "        # shape = [1, batch size]\n",
    "        inputs = inputs.expand_dims(0)\n",
    "\n",
    "        # 输入词的embedding输出， d(y_t)\n",
    "        # shape = [1, batch size, emb dim]\n",
    "        embedded = self.dropout(self.embedding(inputs))\n",
    "\n",
    "        # 注意力权重向量, a_t\n",
    "        # shape = [batch size, src len]\n",
    "        a = self.attention(hidden, encoder_outputs, mask)\n",
    "\n",
    "        # 为注意力权重增加额外维度\n",
    "        # shape = [batch size, 1, src len]\n",
    "        a = a.expand_dims(1)\n",
    "\n",
    "        # 将编码器隐藏状态中的第1、2维度进行交换\n",
    "        # shape = [batch size, src len, enc hid dim * 2]\n",
    "        encoder_outputs = encoder_outputs.transpose(1, 0, 2)\n",
    "\n",
    "        # 计算w_t\n",
    "        # shape = [batch size, 1, enc hid dim * 2]\n",
    "        weighted = ops.bmm(a, encoder_outputs)\n",
    "\n",
    "        # 将w_t的第1、2维度进行交换\n",
    "        # shape = [1, batch size, enc hid dim * 2]\n",
    "        weighted = weighted.transpose(1, 0, 2)\n",
    "\n",
    "        # 将emdedded与weighted堆叠在一起，后输入进RNN层\n",
    "        # rnn_input shape = [1, batch size, (enc hid dim * 2) + emb dim]\n",
    "        # output shape = [seq len = 1, batch size, dec hid dim * n directions]\n",
    "        # hidden shape = [n layers (1) * n directions (1) = 1, batch size, dec hid dim]\n",
    "        rnn_input = ops.concat((embedded, weighted), axis=2)\n",
    "        output, hidden = self.rnn(rnn_input, hidden.expand_dims(0))\n",
    "\n",
    "        # 去除多余的第1维度\n",
    "        embedded = embedded.squeeze(0)\n",
    "        output = output.squeeze(0)\n",
    "        weighted = weighted.squeeze(0)\n",
    "\n",
    "        # 将embedded，weighted和hidden堆叠起来，并输入线性层，预测下一个词\n",
    "        # shape = [batch size, output dim]\n",
    "        prediction = self.fc_out(ops.concat((output, weighted, embedded), axis=1))\n",
    "\n",
    "        return prediction, hidden.squeeze(0), a.squeeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6d8b6d2-8346-49bc-bada-f18c6f557b5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Cell):\n",
    "    def __init__(self, encoder, decoder, src_pad_idx, teacher_forcing_ratio):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.teacher_forcing_ratio = teacher_forcing_ratio  # 使用teacher forcing的可能性\n",
    "\n",
    "    def create_mask(self, src):\n",
    "        \"\"\"标记出每个序列中<pad>占位符的位置\"\"\"\n",
    "        mask = (src != self.src_pad_idx).astype(mindspore.int32).swapaxes(1, 0)\n",
    "        return mask\n",
    "\n",
    "    def construct(self, src, src_len, trg, trg_len=None):\n",
    "        \"\"\"构建seq2seq模型\n",
    "\n",
    "        Args:\n",
    "            src: 源序列；shape = [src len, batch size]\n",
    "            src_len: 源序列长度；shape = [batch size]\n",
    "            trg: 目标序列；shape = [trg len, batch size]\n",
    "            trg_len: 目标序列长度；shape = [trg len, batch size]\n",
    "        \"\"\"\n",
    "        if trg_len is None:\n",
    "            trg_len = trg.shape[0]\n",
    "\n",
    "        #存储解码器输出\n",
    "        outputs = []\n",
    "\n",
    "        # 编码器（encoder）：\n",
    "        # 输入：源序列、源序列长度\n",
    "        # 输出1：编码器中所有前向与反向RNN 的隐藏状态 encoder_outputs\n",
    "        # 输出2：编码器中前向与反向RNN中最后时刻的隐藏状态放入线性层后的输出 hidden\n",
    "        encoder_outputs, hidden = self.encoder(src, src_len)\n",
    "\n",
    "        #解码器的第一个输入是表示序列开始的占位符<bos>\n",
    "        inputs = trg[0]\n",
    "\n",
    "        # 标记源序列中<pad>占位符的位置\n",
    "        # shape = [batch size, src len]\n",
    "        mask = self.create_mask(src)\n",
    "\n",
    "        for t in range(1, trg_len):\n",
    "\n",
    "            # 解码器（decoder）：\n",
    "            # 输入：源句子序列 inputs、前一时刻的隐藏状态 hidden、编码器所有前向与反向RNN的隐藏状态\n",
    "            # 标明每个句子中的<pad>，方便计算注意力权重时忽略该部分\n",
    "            # 输出：预测结果 output、新的隐藏状态 hidden、注意力权重（忽略）\n",
    "            output, hidden, _ = self.decoder(inputs, hidden, encoder_outputs, mask)\n",
    "\n",
    "            # 将预测结果放入之前的存储中\n",
    "            outputs.append(output)\n",
    "\n",
    "            #找出对应预测概率最大的词元\n",
    "            top1 = output.argmax(1).astype(mindspore.int32)\n",
    "\n",
    "            if self.training:\n",
    "                #如果目前为模型训练状态，则按照之前设定的概率使用teacher forcing\n",
    "                teacher_force = random.random() < self.teacher_forcing_ratio\n",
    "                # 如使用teacher forcing，则将目标序列中对应的词元作为下一个输入\n",
    "                # 如不使用teacher forcing，则将预测结果作为下一个输入\n",
    "                inputs = trg[t] if teacher_force else top1\n",
    "            else:\n",
    "                inputs = top1\n",
    "\n",
    "        # 将所有输出整合为tensor\n",
    "        outputs = ops.stack(outputs, axis=0)\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a928f640-6157-412b-b099-293c8a129fdc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_dim = len(de_vocab)  # 输入维度\n",
    "output_dim = len(en_vocab)  # 输出维度\n",
    "enc_emb_dim = 256  # Encoder Embedding层维度\n",
    "dec_emb_dim = 256  # Decoder Embedding层维度\n",
    "enc_hid_dim = 512  # Encoder 隐藏层维度\n",
    "dec_hid_dim = 512  # Decoder 隐藏层维度\n",
    "enc_dropout = 0.5  # Encoder Dropout\n",
    "dec_dropout = 0.5  # Decoder Dropout\n",
    "src_pad_idx = de_vocab('<pad>')  # 德语词典中pad占位符的数字索引\n",
    "trg_pad_idx = en_vocab('<pad>')  # 英语词典中pad占位符的数字索引\n",
    "\n",
    "\n",
    "attn = Attention(enc_hid_dim, dec_hid_dim)\n",
    "encoder = Encoder(input_dim, enc_emb_dim, enc_hid_dim, dec_hid_dim, enc_dropout)\n",
    "decoder = Decoder(output_dim, dec_emb_dim, enc_hid_dim, dec_hid_dim, dec_dropout, attn)\n",
    "\n",
    "model = Seq2Seq(encoder, decoder, src_pad_idx, 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00c8994e-c899-4b1c-ae74-76629d4e0732",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "opt = nn.Adam(model.trainable_params(), learning_rate=0.001)  # 损失函数\n",
    "loss_fn = nn.CrossEntropyLoss(ignore_index=trg_pad_idx)  # 优化器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e802c27-6afd-4303-88bc-3dd78611ef63",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def clip_by_norm(clip_norm, t, axis=None):\n",
    "    \"\"\"给定张量t和裁剪参数clip_norm，对t进行正则化\n",
    "\n",
    "    使得t在axes维度上的L2-norm小于等于clip_norm。\n",
    "\n",
    "    Args:\n",
    "        t: tensor，数据类型为float\n",
    "        clip_norm: scalar，数值需大于0；梯度裁剪阈值，数据类型为float\n",
    "        axis: Union[None, int, tuple(int)]，数据类型为int32；计算L2-norm参考的维度，如为Norm，则参考所有维度\n",
    "    \"\"\"\n",
    "\n",
    "    # 计算L2-norm\n",
    "    t2 = t * t\n",
    "    l2sum = t2.sum(axis=axis, keepdims=True)\n",
    "    pred = l2sum > 0\n",
    "    # 将加和中等于0的元素替换为1，避免后续出现NaN\n",
    "    l2sum_safe = ops.select(pred, l2sum, ops.ones_like(l2sum))\n",
    "    l2norm = ops.select(pred, ops.sqrt(l2sum_safe), l2sum)\n",
    "    # 比较L2-norm和clip_norm，如L2-norm超过阈值，进行裁剪\n",
    "    # 剪裁方法：output(x) = (x * clip_norm)/max(|x|, clip_norm)\n",
    "    intermediate = t * clip_norm\n",
    "    cond = l2norm > clip_norm\n",
    "    t_clip = ops.identity(intermediate / ops.select(cond, l2norm, clip_norm))\n",
    "\n",
    "    return t_clip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f4e37c-8364-4e2c-bbf8-5d6f22bc9364",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def forward_fn(src, src_len, trg):\n",
    "    \"\"\"前向网络\"\"\"\n",
    "    src = src.swapaxes(0, 1)\n",
    "    trg = trg.swapaxes(0, 1)\n",
    "\n",
    "    output = model(src, src_len, trg)\n",
    "    output_dim = output.shape[-1]\n",
    "    output = output.view(-1, output_dim)\n",
    "    trg = trg[1:].view(-1)\n",
    "    loss = loss_fn(output, trg)\n",
    "\n",
    "    return loss\n",
    "\n",
    "# 反向传播计算梯度\n",
    "grad_fn = mindspore.value_and_grad(forward_fn, None, opt.parameters)\n",
    "\n",
    "@mindspore.jit\n",
    "def train_step(src, src_len, trg, clip):\n",
    "    \"\"\"单步训练\"\"\"\n",
    "    loss, grads = grad_fn(src, src_len, trg)\n",
    "    grads = map(partial(clip_by_norm, clip), grads)  # 梯度裁剪\n",
    "    opt(grads)  # 更新网络参数\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def train(dataset, clip, epoch=0):\n",
    "    \"\"\"模型训练\"\"\"\n",
    "    model.set_train(True)\n",
    "    num_batches = dataset.get_dataset_size()\n",
    "    total_loss = 0  # 所有batch训练loss的累加\n",
    "    total_steps = 0  # 训练步数\n",
    "\n",
    "    with tqdm(total=num_batches) as t:\n",
    "        t.set_description(f'Epoch: {epoch}')\n",
    "        for src, src_len, trg in dataset.create_tuple_iterator():\n",
    "            loss = train_step(src, src_len.astype(src.dtype), trg, clip)  # 当前batch的loss\n",
    "            total_loss += loss.asnumpy()\n",
    "            total_steps += 1\n",
    "            curr_loss = total_loss / total_steps  # 当前的平均loss\n",
    "            t.set_postfix({'loss': f'{curr_loss:.2f}'})\n",
    "            t.update(1)\n",
    "\n",
    "    return total_loss / total_steps\n",
    "\n",
    "\n",
    "def evaluate(dataset):\n",
    "    \"\"\"模型验证\"\"\"\n",
    "    model.set_train(False)\n",
    "    num_batches = dataset.get_dataset_size()\n",
    "    total_loss = 0  # 所有batch训练loss的累加\n",
    "    total_steps = 0  # 训练步数\n",
    "\n",
    "    with tqdm(total=num_batches) as t:\n",
    "        for src, src_len, trg in dataset.create_tuple_iterator():\n",
    "            loss = forward_fn(src, src_len, trg)  # 当前batch的loss\n",
    "            total_loss += loss.asnumpy()\n",
    "            total_steps += 1\n",
    "            curr_loss = total_loss / total_steps  # 当前的平均loss\n",
    "            t.set_postfix({'loss': f'{curr_loss:.2f}'})\n",
    "            t.update(1)\n",
    "\n",
    "    return total_loss / total_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a9bf37-fdd5-473f-8cf2-4e6437593f03",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from mindspore import save_checkpoint\n",
    "\n",
    "num_epochs = 10  # 训练迭代数\n",
    "clip = 1.0  # 梯度裁剪阈值\n",
    "best_valid_loss = float('inf')  # 当前最佳验证损失\n",
    "\n",
    "for i in range(num_epochs):\n",
    "    # 模型训练，网络权重更新\n",
    "    train_loss = train(train_dataset, clip, i)\n",
    "    # 网络权重更新后对模型进行验证\n",
    "    valid_loss = evaluate(valid_dataset)\n",
    "\n",
    "    # 保存当前效果最好的模型\n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        save_checkpoint(model, 'seq2seq.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0148be89-2385-46a9-8c5e-2ce5b8de922c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
